// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include "testframe.h"

#include <vespa/documentapi/documentapi.h>
#include <vespa/documentapi/messagebus/policies/andpolicy.h>
#include <vespa/documentapi/messagebus/policies/documentrouteselectorpolicy.h>
#include <vespa/documentapi/messagebus/policies/errorpolicy.h>
#include <vespa/documentapi/messagebus/policies/externpolicy.h>
#include <vespa/documentapi/messagebus/policies/loadbalancerpolicy.h>
#include <vespa/documentapi/messagebus/policies/localservicepolicy.h>
#include <vespa/documentapi/messagebus/policies/roundrobinpolicy.h>
#include <vespa/documentapi/messagebus/policies/searchcolumnpolicy.h>
#include <vespa/documentapi/messagebus/policies/searchrowpolicy.h>
#include <vespa/documentapi/messagebus/policies/storagepolicy.h>
#include <vespa/documentapi/messagebus/policies/subsetservicepolicy.h>
#include <vespa/messagebus/emptyreply.h>
#include <vespa/messagebus/routing/routingnode.h>
#include <vespa/messagebus/routing/routingtable.h>
#include <vespa/messagebus/routing/policydirective.h>
#include <vespa/messagebus/testlib/testserver.h>
#include <vespa/vdslib/container/mutabledocumentlist.h>
#include <vespa/vdslib/state/clusterstate.h>
#include <vespa/document/base/testdocrepo.h>
#include <vespa/document/fieldvalue/longfieldvalue.h>
#include <vespa/document/datatype/documenttype.h>
#include <vespa/vespalib/testkit/testapp.h>
#include <vespa/vespalib/util/stringfmt.h>

#include <vespa/log/log.h>
LOG_SETUP("policies_test");

using document::DataType;
using document::Document;
using document::DocumentId;
using document::DocumentTypeRepo;
using document::DocumentUpdate;
using document::readDocumenttypesConfig;
using slobrok::api::IMirrorAPI;
using namespace documentapi;

class Test : public vespalib::TestApp {
private:
    LoadTypeSet          _loadTypes;
    DocumentTypeRepo::SP _repo;
    const DataType      *_docType;

private:
    bool trySelect(TestFrame &frame, uint32_t numSelects, const std::vector<string> &expected);
    bool tryDistribution(TestFrame &frame, const string &id, const string &expected);
    void tryWasFound(TestFrame &frame, uint32_t expectedRecipients,
                     uint32_t foundMask, bool expectedFound);
    void setupExternPolicy(TestFrame &frame, mbus::Slobrok &slobrok, const string &pattern,
                           int32_t numEntries = -1);
    StoragePolicy &setupStoragePolicy(TestFrame &frame, const string &param,
                                      const string &pattern = "", int32_t numEntries = -1);
    bool isErrorPolicy(const string &name, const string &param);
    void assertMirrorReady(const IMirrorAPI &mirror);
    void assertMirrorContains(const IMirrorAPI &mirror, const string &pattern,
                              uint32_t numEntries);
    mbus::Message::UP newPutDocumentMessage(const string &documentId);

public:
    Test();
    ~Test();
    int  Main() override;
    void testAND();
    void testDocumentRouteSelector();
    void testDocumentRouteSelectorIgnore();
    void testExternSend();
    void testExternMultipleSlobroks();
    void testLoadBalancer();
    void testLocalService();
    void testLocalServiceCache();
    void testProtocol();
    void testRoundRobin();
    void testRoundRobinCache();
    void testSearchColumn();
    void testSearchRow();
    void testSearchRowMerge();
    void multipleGetRepliesAreMergedToFoundDocument();
    void testSubsetService();
    void testSubsetServiceCache();

    void requireThatExternPolicyWithIllegalParamIsAnErrorPolicy();
    void requireThatExternPolicyWithUnknownPatternSelectsNone();
    void requireThatExternPolicySelectsFromExternSlobrok();
    void requireThatExternPolicyMergesOneReplyAsProtocol();
    void requireThatStoragePolicyWithIllegalParamIsAnErrorPolicy();
    void requireThatStoragePolicyIsRandomWithoutState();
    void requireThatStoragePolicyIsTargetedWithState();
    void requireThatStoragePolicyCombinesSystemAndSlobrokState();
};

TEST_APPHOOK(Test);

Test::Test() {}
Test::~Test() {}

int
Test::Main() {
    TEST_INIT(_argv[0]);

    _repo.reset(new DocumentTypeRepo(readDocumenttypesConfig(
                TEST_PATH("../../../test/cfg/testdoctypes.cfg"))));
    _docType = _repo->getDocumentType("testdoc");

    testProtocol();                     TEST_FLUSH();

    testAND();                          TEST_FLUSH();
    testDocumentRouteSelector();        TEST_FLUSH();
    testDocumentRouteSelectorIgnore();  TEST_FLUSH();
    testExternSend();                   TEST_FLUSH();
    testExternMultipleSlobroks();       TEST_FLUSH();
    testLoadBalancer();                 TEST_FLUSH();
    testLocalService();                 TEST_FLUSH();
    testLocalServiceCache();            TEST_FLUSH();
    testRoundRobin();                   TEST_FLUSH();
    testRoundRobinCache();              TEST_FLUSH();
    testSearchColumn();                 TEST_FLUSH();
    testSearchRow();                    TEST_FLUSH();
    testSearchRowMerge();               TEST_FLUSH();
    testSubsetService();                TEST_FLUSH();
    testSubsetServiceCache();           TEST_FLUSH();

    multipleGetRepliesAreMergedToFoundDocument(); TEST_FLUSH();

    requireThatExternPolicyWithIllegalParamIsAnErrorPolicy();  TEST_FLUSH();
    requireThatExternPolicyWithUnknownPatternSelectsNone();    TEST_FLUSH();
    requireThatExternPolicySelectsFromExternSlobrok();         TEST_FLUSH();
    requireThatExternPolicyMergesOneReplyAsProtocol();         TEST_FLUSH();

    requireThatStoragePolicyWithIllegalParamIsAnErrorPolicy(); TEST_FLUSH();
    requireThatStoragePolicyIsRandomWithoutState();            TEST_FLUSH();
    requireThatStoragePolicyIsTargetedWithState();             TEST_FLUSH();
    requireThatStoragePolicyCombinesSystemAndSlobrokState();   TEST_FLUSH();

    TEST_DONE();
}

void
Test::testProtocol()
{
    mbus::IProtocol::SP protocol(new DocumentProtocol(_loadTypes, _repo));

    mbus::IRoutingPolicy::UP policy = protocol->createPolicy("AND", "");
    ASSERT_TRUE(dynamic_cast<ANDPolicy*>(policy.get()) != NULL);

    policy = protocol->createPolicy("DocumentRouteSelector", "raw:route[0]\n");
    ASSERT_TRUE(dynamic_cast<DocumentRouteSelectorPolicy*>(policy.get()) != NULL);

    policy = protocol->createPolicy("Extern", "foo;bar/baz");
    ASSERT_TRUE(dynamic_cast<ExternPolicy*>(policy.get()) != NULL);

    policy = protocol->createPolicy("LoadBalancer",
                                    "cluster=docproc/cluster.default;"
                                    "session=chain.default;syncinit");
    ASSERT_TRUE(dynamic_cast<LoadBalancerPolicy*>(policy.get()) != NULL);

    policy = protocol->createPolicy("LocalService", "");
    ASSERT_TRUE(dynamic_cast<LocalServicePolicy*>(policy.get()) != NULL);

    policy = protocol->createPolicy("RoundRobin", "");
    ASSERT_TRUE(dynamic_cast<RoundRobinPolicy*>(policy.get()) != NULL);

    policy = protocol->createPolicy("SearchRow", "");
    ASSERT_TRUE(dynamic_cast<SearchRowPolicy*>(policy.get()) != NULL);

    policy = protocol->createPolicy("SearchColumn", "");
    ASSERT_TRUE(dynamic_cast<SearchColumnPolicy*>(policy.get()) != NULL);

    policy = protocol->createPolicy("SubsetService", "");
    ASSERT_TRUE(dynamic_cast<SubsetServicePolicy*>(policy.get()) != NULL);
}

void
Test::testAND()
{
    TestFrame frame(_repo);
    frame.setMessage(mbus::Message::UP(new PutDocumentMessage(
                            document::Document::SP(
                                    new document::Document(*_docType,
                                            DocumentId("doc:scheme:"))))));
    frame.setHop(mbus::HopSpec("test", "[AND]")
                 .addRecipient("foo")
                 .addRecipient("bar"));
    EXPECT_TRUE(frame.testSelect(StringList().add("foo").add("bar")));

    frame.setHop(mbus::HopSpec("test", "[AND:baz]")
                 .addRecipient("foo").addRecipient("bar"));
    EXPECT_TRUE(frame.testSelect(StringList().add("baz"))); // param precedes recipients

    frame.setHop(mbus::HopSpec("test", "[AND:foo]"));
    EXPECT_TRUE(frame.testMergeOneReply("foo"));

    frame.setHop(mbus::HopSpec("test", "[AND:foo bar]"));
    EXPECT_TRUE(frame.testMergeTwoReplies("foo", "bar"));
}

void
Test::requireThatExternPolicyWithIllegalParamIsAnErrorPolicy()
{
    mbus::Slobrok slobrok;
    string spec = vespalib::make_string("tcp/localhost:%d", slobrok.port());

    EXPECT_TRUE(isErrorPolicy("Extern", ""));
    EXPECT_TRUE(isErrorPolicy("Extern", spec));
    EXPECT_TRUE(isErrorPolicy("Extern", spec + ";"));
    EXPECT_TRUE(isErrorPolicy("Extern", spec + ";bar"));
}

void
Test::requireThatExternPolicyWithUnknownPatternSelectsNone()
{
    TestFrame frame(_repo);
    frame.setMessage(newPutDocumentMessage("doc:scheme:"));

    mbus::Slobrok slobrok;
    setupExternPolicy(frame, slobrok, "foo/bar");
    EXPECT_TRUE(frame.testSelect(StringList()));
}

void
Test::requireThatExternPolicySelectsFromExternSlobrok()
{
    TestFrame frame(_repo);
    frame.setMessage(newPutDocumentMessage("doc:scheme:"));
    mbus::Slobrok slobrok;
    std::vector<mbus::TestServer*> servers;
    for (uint32_t i = 0; i < 10; ++i) {
        mbus::TestServer *server = new mbus::TestServer(
                mbus::Identity(vespalib::make_string("docproc/cluster.default/%d", i)),
                mbus::RoutingSpec(), slobrok, "",
                mbus::IProtocol::SP(new DocumentProtocol(_loadTypes, _repo)));
        servers.push_back(server);
        server->net.registerSession("chain.default");
    }
    setupExternPolicy(frame, slobrok, "docproc/cluster.default/*/chain.default", 10);
    std::set<string> lst;
    for (uint32_t i = 0; i < servers.size(); ++i) {
        std::vector<mbus::RoutingNode*> leaf;
        ASSERT_TRUE(frame.select(leaf, 1));
        lst.insert(leaf[0]->getRoute().toString());

        leaf[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));
        ASSERT_TRUE(frame.getReceptor().getReply(600).get() != NULL);
    }
    EXPECT_EQUAL(servers.size(), lst.size());
    for (uint32_t i = 0; i < servers.size(); ++i) {
        delete servers[i];
    }
}

void
Test::requireThatExternPolicyMergesOneReplyAsProtocol()
{
    TestFrame frame(_repo);
    frame.setMessage(newPutDocumentMessage("doc:scheme:"));
    mbus::Slobrok slobrok;
    mbus::TestServer server(mbus::Identity("docproc/cluster.default/0"),
                            mbus::RoutingSpec(), slobrok, "",
                            mbus::IProtocol::SP(new DocumentProtocol(_loadTypes, _repo)));
    server.net.registerSession("chain.default");
    setupExternPolicy(frame, slobrok, "docproc/cluster.default/0/chain.default", 1);
    EXPECT_TRUE(frame.testMergeOneReply(server.net.getConnectionSpec() + "/chain.default"));
}

mbus::Message::UP
Test::newPutDocumentMessage(const string &documentId)
{
    Document::SP doc(new Document(*_docType, DocumentId(documentId)));
    return mbus::Message::UP(new PutDocumentMessage(doc));
}

void
Test::setupExternPolicy(TestFrame &frame, mbus::Slobrok &slobrok, const string &pattern,
                        int32_t numEntries)
{
    string param = vespalib::make_string("tcp/localhost:%d;%s",
                                         slobrok.port(), pattern.c_str());
    frame.setHop(mbus::HopSpec("test", vespalib::make_string("[Extern:%s]", param.c_str())));
    mbus::MessageBus &mbus = frame.getMessageBus();
    const mbus::HopBlueprint *hop = mbus.getRoutingTable(DocumentProtocol::NAME)->getHop("test");
    const mbus::PolicyDirective dir = static_cast<mbus::PolicyDirective&>(*hop->getDirective(0));
    ExternPolicy &policy = static_cast<ExternPolicy&>(*mbus.getRoutingPolicy(
                    DocumentProtocol::NAME,
                    dir.getName(),
                    dir.getParam()));
    assertMirrorReady(policy.getMirror());
    if (numEntries >= 0) {
        assertMirrorContains(policy.getMirror(), pattern, numEntries);
    }
}

void
Test::assertMirrorReady(const slobrok::api::IMirrorAPI &mirror)
{
    for (uint32_t i = 0; i < 6000; ++i) {
        if (mirror.ready()) {
            return;
        }
        FastOS_Thread::Sleep(10);
    }
    ASSERT_TRUE(false);
}

void
Test::assertMirrorContains(const slobrok::api::IMirrorAPI &mirror, const string &pattern,
                           uint32_t numEntries)
{
    for (uint32_t i = 0; i < 6000; ++i) {
        if (mirror.lookup(pattern).size() == numEntries) {
            return;
        }
        FastOS_Thread::Sleep(10);
    }
    ASSERT_TRUE(false);
}

void
Test::testExternSend()
{
    // Setup local source node.
    mbus::Slobrok local;
    mbus::TestServer src(mbus::Identity("src"), mbus::RoutingSpec(), local, "",
                         mbus::IProtocol::SP(
                                 new DocumentProtocol(_loadTypes, _repo)));
    mbus::Receptor sr;
    mbus::SourceSession::UP ss = src.mb.createSourceSession(sr, mbus::SourceSessionParams().setTimeout(60));

    mbus::Slobrok slobrok;
    mbus::TestServer itr(mbus::Identity("itr"), mbus::RoutingSpec()
                         .addTable(mbus::RoutingTableSpec(DocumentProtocol::NAME)
                                   .addRoute(mbus::RouteSpec("default").addHop("dst"))
                                   .addHop(mbus::HopSpec("dst", "dst/session"))),
                         slobrok, "", mbus::IProtocol::SP(
                                 new DocumentProtocol(_loadTypes, _repo)));
    mbus::Receptor ir;
    mbus::IntermediateSession::UP is = itr.mb.createIntermediateSession("session", true, ir, ir);

    mbus::TestServer dst(mbus::Identity("dst"), mbus::RoutingSpec(), slobrok, "",
                         mbus::IProtocol::SP(new DocumentProtocol(_loadTypes, _repo)));
    mbus::Receptor dr;
    mbus::DestinationSession::UP ds = dst.mb.createDestinationSession("session", true, dr);

    // Send message from local node to remote cluster and resolve route there.
    mbus::Message::UP msg(new GetDocumentMessage(document::DocumentId("doc:scheme:"), 0));
    msg->getTrace().setLevel(9);
    msg->setRoute(mbus::Route::parse(vespalib::make_string("[Extern:tcp/localhost:%d;itr/session] default", slobrok.port())));

    ASSERT_TRUE(ss->send(std::move(msg)).isAccepted());
    ASSERT_TRUE((msg = ir.getMessage(600)).get() != NULL);
    is->forward(std::move(msg));
    ASSERT_TRUE((msg = dr.getMessage(600)).get() != NULL);
    ds->acknowledge(std::move(msg));
    mbus::Reply::UP reply = ir.getReply(600);
    ASSERT_TRUE(reply.get() != NULL);
    is->forward(std::move(reply));
    ASSERT_TRUE((reply = sr.getReply(600)).get() != NULL);

    fprintf(stderr, "%s", reply->getTrace().toString().c_str());
}

void
Test::testExternMultipleSlobroks()
{
    mbus::Slobrok local;
    mbus::TestServer src(mbus::Identity("src"), mbus::RoutingSpec(), local, "",
                         mbus::IProtocol::SP(new DocumentProtocol(_loadTypes, _repo)));
    mbus::Receptor sr;
    mbus::SourceSession::UP ss = src.mb.createSourceSession(sr, mbus::SourceSessionParams().setTimeout(60));

    string spec;
    mbus::Receptor dr;
    {
        mbus::Slobrok ext;
        spec.append(vespalib::make_string("tcp/localhost:%d", ext.port()));

        mbus::TestServer dst(mbus::Identity("dst"), mbus::RoutingSpec(), ext, "",
                             mbus::IProtocol::SP(new DocumentProtocol(_loadTypes, _repo)));
        mbus::DestinationSession::UP ds = dst.mb.createDestinationSession("session", true, dr);

        mbus::Message::UP msg(new GetDocumentMessage(document::DocumentId("doc:scheme:"), 0));
        msg->setRoute(mbus::Route::parse(vespalib::make_string("[Extern:%s;dst/session]", spec.c_str())));
        ASSERT_TRUE(ss->send(std::move(msg)).isAccepted());
        ASSERT_TRUE((msg = dr.getMessage(600)).get() != NULL);
        ds->acknowledge(std::move(msg));
        mbus::Reply::UP reply = sr.getReply(600);
        ASSERT_TRUE(reply.get() != NULL);
    }
    {
        mbus::Slobrok ext;
        spec.append(vespalib::make_string(",tcp/localhost:%d", ext.port()));

        mbus::TestServer dst(mbus::Identity("dst"), mbus::RoutingSpec(), ext, "",
                             mbus::IProtocol::SP(new DocumentProtocol(_loadTypes, _repo)));
        mbus::DestinationSession::UP ds = dst.mb.createDestinationSession("session", true, dr);

        mbus::Message::UP msg(new GetDocumentMessage(document::DocumentId("doc:scheme:"), 0));
        msg->setRoute(mbus::Route::parse(vespalib::make_string("[Extern:%s;dst/session]", spec.c_str())));
        ASSERT_TRUE(ss->send(std::move(msg)).isAccepted());
        ASSERT_TRUE((msg = dr.getMessage(600)).get() != NULL);
        ds->acknowledge(std::move(msg));
        mbus::Reply::UP reply = sr.getReply(600);
        ASSERT_TRUE(reply.get() != NULL);
    }
}

void
Test::testLocalService()
{
    // Prepare message.
    TestFrame frame(_repo, "docproc/cluster.default");
    frame.setMessage(mbus::Message::UP(new PutDocumentMessage(Document::SP(
                                    new Document(*_docType,
                                            DocumentId("doc:scheme:"))))));

    // Test select with proper address.
    for (uint32_t i = 0; i < 10; ++i) {
        frame.getNetwork().registerSession(vespalib::make_string("%d/chain.default", i));
    }
    ASSERT_TRUE(frame.waitSlobrok("docproc/cluster.default/*/chain.default", 10));
    frame.setHop(mbus::HopSpec("test", "docproc/cluster.default/[LocalService]/chain.default"));

    std::set<string> lst;
    for (uint32_t i = 0; i < 10; ++i) {
        std::vector<mbus::RoutingNode*> leaf;
        ASSERT_TRUE(frame.select(leaf, 1));
        lst.insert(leaf[0]->getRoute().toString());

        leaf[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));
        ASSERT_TRUE(frame.getReceptor().getReply(600).get() != NULL);
    }
    EXPECT_EQUAL(10u, lst.size());

    // Test select with broken address.
    lst.clear();
    frame.setHop(mbus::HopSpec("test", "docproc/cluster.default/[LocalService:broken]/chain.default"));
    for (uint32_t i = 0; i < 10; ++i) {
        std::vector<mbus::RoutingNode*> leaf;
        ASSERT_TRUE(frame.select(leaf, 1));
        lst.insert(leaf[0]->getRoute().toString());

        leaf[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));
        ASSERT_TRUE(frame.getReceptor().getReply(600).get() != NULL);
    }
    EXPECT_EQUAL(1u, lst.size());
    EXPECT_EQUAL("docproc/cluster.default/*/chain.default", *lst.begin());

    // Test merge behavior.
    frame.setHop(mbus::HopSpec("test", "[LocalService]"));
    EXPECT_TRUE(frame.testMergeOneReply("*"));
}

void
Test::testLocalServiceCache()
{
    TestFrame fooFrame(_repo, "docproc/cluster.default");
    mbus::HopSpec fooHop("foo", "docproc/cluster.default/[LocalService]/chain.foo");
    fooFrame.setMessage(mbus::Message::UP(new GetDocumentMessage(document::DocumentId("doc:scheme:foo"))));
    fooFrame.setHop(fooHop);

    TestFrame barFrame(fooFrame);
    mbus::HopSpec barHop("test", "docproc/cluster.default/[LocalService]/chain.bar");
    barFrame.setMessage(mbus::Message::UP(new GetDocumentMessage(document::DocumentId("doc:scheme:bar"))));
    barFrame.setHop(barHop);

    fooFrame.getMessageBus().setupRouting(
            mbus::RoutingSpec().addTable(mbus::RoutingTableSpec(DocumentProtocol::NAME)
                                         .addHop(fooHop)
                                         .addHop(barHop)));

    fooFrame.getNetwork().registerSession("0/chain.foo");
    fooFrame.getNetwork().registerSession("0/chain.bar");
    ASSERT_TRUE(fooFrame.waitSlobrok("docproc/cluster.default/0/*", 2));

    std::vector<mbus::RoutingNode*> fooSelected;
    ASSERT_TRUE(fooFrame.select(fooSelected, 1));
    EXPECT_EQUAL("docproc/cluster.default/0/chain.foo", fooSelected[0]->getRoute().getHop(0).toString());

    std::vector<mbus::RoutingNode*> barSelected;
    ASSERT_TRUE(barFrame.select(barSelected, 1));
    EXPECT_EQUAL("docproc/cluster.default/0/chain.bar", barSelected[0]->getRoute().getHop(0).toString());

    barSelected[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));
    fooSelected[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));

    ASSERT_TRUE(barFrame.getReceptor().getReply(600).get() != NULL);
    ASSERT_TRUE(fooFrame.getReceptor().getReply(600).get() != NULL);
}

void
Test::testRoundRobin()
{
    // Prepare message.
    TestFrame frame(_repo, "docproc/cluster.default");
    frame.setMessage(mbus::Message::UP(new PutDocumentMessage(Document::SP(
                                    new Document(*_docType,
                                            DocumentId("doc:scheme:"))))));

    // Test select with proper address.
    for (uint32_t i = 0; i < 10; ++i) {
        frame.getNetwork().registerSession(vespalib::make_string("%d/chain.default", i));
    }
    ASSERT_TRUE(frame.waitSlobrok("docproc/cluster.default/*/chain.default", 10));
    frame.setHop(mbus::HopSpec("test", "[RoundRobin]")
                 .addRecipient("docproc/cluster.default/3/chain.default")
                 .addRecipient("docproc/cluster.default/6/chain.default")
                 .addRecipient("docproc/cluster.default/9/chain.default"));
    EXPECT_TRUE(trySelect(frame, 32, StringList()
                          .add("docproc/cluster.default/3/chain.default")
                          .add("docproc/cluster.default/6/chain.default")
                          .add("docproc/cluster.default/9/chain.default")));
    frame.getNetwork().unregisterSession("6/chain.default");
    ASSERT_TRUE(frame.waitSlobrok("docproc/cluster.default/*/chain.default", 9));
    EXPECT_TRUE(trySelect(frame, 32, StringList()
                          .add("docproc/cluster.default/3/chain.default")
                          .add("docproc/cluster.default/9/chain.default")));
    frame.getNetwork().unregisterSession("3/chain.default");
    ASSERT_TRUE(frame.waitSlobrok("docproc/cluster.default/*/chain.default", 8));
    EXPECT_TRUE(trySelect(frame, 32, StringList()
                          .add("docproc/cluster.default/9/chain.default")));
    frame.getNetwork().unregisterSession("9/chain.default");
    ASSERT_TRUE(frame.waitSlobrok("docproc/cluster.default/*/chain.default", 7));
    EXPECT_TRUE(trySelect(frame, 32, StringList()));

    // Test merge behavior.
    frame.setHop(mbus::HopSpec("test", "[RoundRobin]").addRecipient("docproc/cluster.default/0/chain.default"));
    EXPECT_TRUE(frame.testMergeOneReply("docproc/cluster.default/0/chain.default"));
}

void
Test::testRoundRobinCache()
{
    TestFrame fooFrame(_repo, "docproc/cluster.default");
    mbus::HopSpec fooHop("foo", "[RoundRobin]");
    fooHop.addRecipient("docproc/cluster.default/0/chain.foo");
    fooFrame.setMessage(mbus::Message::UP(new GetDocumentMessage(document::DocumentId("doc:scheme:foo"))));
    fooFrame.setHop(fooHop);

    TestFrame barFrame(fooFrame);
    mbus::HopSpec barHop("bar", "[RoundRobin]");
    barHop.addRecipient("docproc/cluster.default/0/chain.bar");
    barFrame.setMessage(mbus::Message::UP(new GetDocumentMessage(document::DocumentId("doc:scheme:bar"))));
    barFrame.setHop(barHop);

    fooFrame.getMessageBus().setupRouting(
            mbus::RoutingSpec().addTable(mbus::RoutingTableSpec(DocumentProtocol::NAME)
                                         .addHop(fooHop)
                                         .addHop(barHop)));

    fooFrame.getNetwork().registerSession("0/chain.foo");
    fooFrame.getNetwork().registerSession("0/chain.bar");
    ASSERT_TRUE(fooFrame.waitSlobrok("docproc/cluster.default/0/*", 2));

    std::vector<mbus::RoutingNode*> fooSelected;
    ASSERT_TRUE(fooFrame.select(fooSelected, 1));
    EXPECT_EQUAL("docproc/cluster.default/0/chain.foo", fooSelected[0]->getRoute().getHop(0).toString());

    std::vector<mbus::RoutingNode*> barSelected;
    ASSERT_TRUE(barFrame.select(barSelected, 1));
    EXPECT_EQUAL("docproc/cluster.default/0/chain.bar", barSelected[0]->getRoute().getHop(0).toString());

    barSelected[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));
    fooSelected[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));

    ASSERT_TRUE(barFrame.getReceptor().getReply(600).get() != NULL);
    ASSERT_TRUE(fooFrame.getReceptor().getReply(600).get() != NULL);
}

void
Test::testSearchRow()
{
    TestFrame frame(_repo);
    frame.setMessage(mbus::Message::UP(new PutDocumentMessage(Document::SP(
                                    new Document(*_docType,
                                            DocumentId("doc:scheme:"))))));
    frame.setHop(mbus::HopSpec("test", "[SearchRow]")
                 .addRecipient("foo"));
    EXPECT_TRUE(frame.testMergeOneReply("foo"));
    frame.setHop(mbus::HopSpec("test", "[SearchRow]")
                 .addRecipient("foo")
                 .addRecipient("bar"));
    EXPECT_TRUE(frame.testMergeTwoReplies("foo", "bar"));

    frame.setHop(mbus::HopSpec("test", "[SearchRow:1]")
                 .addRecipient("foo"));
    TestFrame::ReplyMap replies;
    replies["foo"] = mbus::ErrorCode::SERVICE_OOS;
    EXPECT_TRUE(frame.testMergeError(replies, UIntList().add(mbus::ErrorCode::SERVICE_OOS)));

    frame.setHop(mbus::HopSpec("test", "[SearchRow:1]")
                 .addRecipient("foo")
                 .addRecipient("bar"));
    replies["foo"] = mbus::ErrorCode::SERVICE_OOS;
    replies["bar"] = mbus::ErrorCode::NONE;
    EXPECT_TRUE(frame.testMergeOk(replies, StringList().add("bar")));

    replies["foo"] = mbus::ErrorCode::SERVICE_OOS;
    replies["bar"] = mbus::ErrorCode::SERVICE_OOS;
    EXPECT_TRUE(frame.testMergeError(replies, UIntList()
                                    .add(mbus::ErrorCode::SERVICE_OOS)
                                    .add(mbus::ErrorCode::SERVICE_OOS)));

    frame.setHop(mbus::HopSpec("test", "[SearchRow:1]")
                 .addRecipient("foo")
                 .addRecipient("bar")
                 .addRecipient("baz"));
    replies["foo"] = mbus::ErrorCode::SERVICE_OOS;
    replies["bar"] = mbus::ErrorCode::NONE;
    replies["baz"] = mbus::ErrorCode::NONE;
    EXPECT_TRUE(frame.testMergeOk(replies, StringList().add("bar").add("baz")));

    replies["foo"] = mbus::ErrorCode::SERVICE_OOS;
    replies["bar"] = mbus::ErrorCode::SERVICE_OOS;
    replies["baz"] = mbus::ErrorCode::NONE;
    EXPECT_TRUE(frame.testMergeOk(replies, StringList().add("baz")));

    replies["foo"] = mbus::ErrorCode::SERVICE_OOS;
    replies["bar"] = mbus::ErrorCode::SERVICE_OOS;
    replies["baz"] = mbus::ErrorCode::SERVICE_OOS;
    EXPECT_TRUE(frame.testMergeError(replies, UIntList()
                                    .add(mbus::ErrorCode::SERVICE_OOS)
                                    .add(mbus::ErrorCode::SERVICE_OOS)
                                    .add(mbus::ErrorCode::SERVICE_OOS)));

    frame.setHop(mbus::HopSpec("test", "[SearchRow:2]")
                 .addRecipient("foo")
                 .addRecipient("bar")
                 .addRecipient("baz"));
    replies["foo"] = mbus::ErrorCode::SERVICE_OOS;
    replies["bar"] = mbus::ErrorCode::NONE;
    replies["baz"] = mbus::ErrorCode::NONE;
    EXPECT_TRUE(frame.testMergeOk(replies, StringList().add("bar").add("baz")));

    replies["foo"] = mbus::ErrorCode::SERVICE_OOS;
    replies["bar"] = mbus::ErrorCode::SERVICE_OOS;
    replies["baz"] = mbus::ErrorCode::NONE;
    EXPECT_TRUE(frame.testMergeError(replies, UIntList()
                                    .add(mbus::ErrorCode::SERVICE_OOS)
                                    .add(mbus::ErrorCode::SERVICE_OOS)));

    replies["foo"] = mbus::ErrorCode::SERVICE_OOS;
    replies["bar"] = mbus::ErrorCode::SERVICE_OOS;
    replies["baz"] = mbus::ErrorCode::SERVICE_OOS;
    EXPECT_TRUE(frame.testMergeError(replies, UIntList()
                                    .add(mbus::ErrorCode::SERVICE_OOS)
                                    .add(mbus::ErrorCode::SERVICE_OOS)
                                    .add(mbus::ErrorCode::SERVICE_OOS)));
}

void
Test::testSearchRowMerge()
{
    TestFrame frame(_repo);
    frame.setHop(mbus::HopSpec("test", "[SearchRow]")
                 .addRecipient("foo"));
    tryWasFound(frame, 1, 0x0, false);
    tryWasFound(frame, 1, 0x1, true);

    frame.setHop(mbus::HopSpec("test", "[SearchRow]")
                 .addRecipient("foo")
                 .addRecipient("bar"));
    tryWasFound(frame, 2, 0x0, false);
    tryWasFound(frame, 2, 0x1, true);
    tryWasFound(frame, 2, 0x2, true);
    tryWasFound(frame, 2, 0x3, true);

    frame.setHop(mbus::HopSpec("test", "[SearchRow]")
                 .addRecipient("foo")
                 .addRecipient("bar")
                 .addRecipient("baz"));
    tryWasFound(frame, 3, 0x0, false);
    tryWasFound(frame, 3, 0x1, true);
    tryWasFound(frame, 3, 0x2, true);
    tryWasFound(frame, 3, 0x3, true);
    tryWasFound(frame, 3, 0x4, true);
    tryWasFound(frame, 3, 0x5, true);
    tryWasFound(frame, 3, 0x6, true);
    tryWasFound(frame, 3, 0x7, true);
}

void
Test::tryWasFound(TestFrame &frame, uint32_t expectedRecipients,
                  uint32_t foundMask, bool expectedFound)
{
    {
        frame.setMessage(mbus::Message::UP(new RemoveDocumentMessage(DocumentId("doc:scheme:69"))));
        std::vector<mbus::RoutingNode*> selected;
        EXPECT_TRUE(frame.select(selected, expectedRecipients));
        for (uint32_t i = 0, len = selected.size(); i < len; ++i) {
            mbus::Reply::UP reply(new RemoveDocumentReply());
            static_cast<RemoveDocumentReply&>(*reply).setWasFound((1 << i) & foundMask);
            selected[i]->handleReply(std::move(reply));
        }
        mbus::Reply::UP reply = frame.getReceptor().getReply(600);
        EXPECT_TRUE(reply.get() != NULL);
        EXPECT_EQUAL((uint32_t)DocumentProtocol::REPLY_REMOVEDOCUMENT, reply->getType());
        EXPECT_EQUAL(expectedFound, static_cast<RemoveDocumentReply&>(*reply).wasFound());
    }
    {
        DocumentUpdate::SP upd(new DocumentUpdate(*_docType, DocumentId("doc:scheme:")));
        frame.setMessage(mbus::Message::UP(new UpdateDocumentMessage(upd)));
        std::vector<mbus::RoutingNode*> selected;
        EXPECT_TRUE(frame.select(selected, expectedRecipients));
        for (uint32_t i = 0, len = selected.size(); i < len; ++i) {
            mbus::Reply::UP reply(new UpdateDocumentReply());
            static_cast<UpdateDocumentReply&>(*reply).setWasFound((1 << i) & foundMask);
            selected[i]->handleReply(std::move(reply));
        }
        mbus::Reply::UP reply = frame.getReceptor().getReply(600);
        EXPECT_TRUE(reply.get() != NULL);
        EXPECT_EQUAL((uint32_t)DocumentProtocol::REPLY_UPDATEDOCUMENT, reply->getType());
        EXPECT_EQUAL(expectedFound, static_cast<UpdateDocumentReply&>(*reply).wasFound());
    }
}

void
Test::multipleGetRepliesAreMergedToFoundDocument()
{
    TestFrame frame(_repo);
    frame.setHop(mbus::HopSpec("test", "[DocumentRouteSelector:raw:"
                               "route[2]\n"
                               "route[0].name \"foo\"\n"
                               "route[0].selector \"testdoc\"\n"
                               "route[0].feed \"myfeed\"\n"
                               "route[1].name \"bar\"\n"
                               "route[1].selector \"other\"\n"
                               "route[1].feed \"myfeed\"\n]")
                 .addRecipient("foo")
                 .addRecipient("bar"));
    frame.setMessage(mbus::Message::UP(new GetDocumentMessage(DocumentId("doc:scheme:yarn"))));
    std::vector<mbus::RoutingNode*> selected;
    EXPECT_TRUE(frame.select(selected, 2));
    for (uint32_t i = 0, len = selected.size(); i < len; ++i) {
        document::Document::SP doc;
        if (i == 0) {
            doc.reset(new Document(*_docType, DocumentId("doc:scheme:yarn")));
            doc->setLastModified(123456ULL);
        }
        mbus::Reply::UP reply(new GetDocumentReply(std::move(doc)));
        selected[i]->handleReply(std::move(reply));
    }
    mbus::Reply::UP reply = frame.getReceptor().getReply(600);
    EXPECT_TRUE(reply.get() != NULL);
    EXPECT_EQUAL(static_cast<uint32_t>(DocumentProtocol::REPLY_GETDOCUMENT),
                 reply->getType());
    EXPECT_EQUAL(123456ULL, static_cast<GetDocumentReply&>(*reply).getLastModified());
}

void
Test::testSearchColumn()
{
    TestFrame frame(_repo);
    frame.setHop(mbus::HopSpec("test", "[SearchColumn]")
                 .addRecipient("c0")
                 .addRecipient("c1")
                 .addRecipient("c2")
                 .addRecipient("c3"));

    // Test hash distribution.
    EXPECT_TRUE(tryDistribution(frame, "doc:ns:3", "c0"));
    EXPECT_TRUE(tryDistribution(frame, "doc:ns:18", "c1"));
    EXPECT_TRUE(tryDistribution(frame, "doc:ns:0", "c2"));
    EXPECT_TRUE(tryDistribution(frame, "doc:ns:4", "c3"));

    EXPECT_TRUE(tryDistribution(frame, "userdoc:ns:49152:0", "c0"));
    EXPECT_TRUE(tryDistribution(frame, "userdoc:ns:49152:1", "c0"));
    EXPECT_TRUE(tryDistribution(frame, "userdoc:ns:16384:2", "c1"));
    EXPECT_TRUE(tryDistribution(frame, "userdoc:ns:16384:3", "c1"));
    EXPECT_TRUE(tryDistribution(frame, "userdoc:ns:5461:4", "c2"));
    EXPECT_TRUE(tryDistribution(frame, "userdoc:ns:5461:5", "c2"));
    EXPECT_TRUE(tryDistribution(frame, "userdoc:ns:0:6", "c3"));
    EXPECT_TRUE(tryDistribution(frame, "userdoc:ns:0:7", "c3"));

    EXPECT_TRUE(tryDistribution(frame, "groupdoc:ns:0:0", "c0"));
    EXPECT_TRUE(tryDistribution(frame, "groupdoc:ns:0:1", "c0"));
    EXPECT_TRUE(tryDistribution(frame, "groupdoc:ns:4:2", "c1"));
    EXPECT_TRUE(tryDistribution(frame, "groupdoc:ns:4:3", "c1"));
    EXPECT_TRUE(tryDistribution(frame, "groupdoc:ns:2:4", "c2"));
    EXPECT_TRUE(tryDistribution(frame, "groupdoc:ns:2:5", "c2"));
    EXPECT_TRUE(tryDistribution(frame, "groupdoc:ns:7:6", "c3"));
    EXPECT_TRUE(tryDistribution(frame, "groupdoc:ns:7:7", "c3"));

    // Test routing based on message type.
    mbus::Message::UP put(new PutDocumentMessage(Document::SP(
                            new Document(*_docType,
                                    DocumentId("doc:scheme:")))));
}

bool
Test::tryDistribution(TestFrame &frame, const string &id, const string &expected)
{
    Document::SP doc(new Document(*_docType, DocumentId(id)));
    mbus::Message::UP msg(new PutDocumentMessage(doc));
    frame.setMessage(std::move(msg));
    return frame.testSelect(StringList().add(expected));
}

void
Test::testDocumentRouteSelector()
{
    // Test policy with usage safeguard.
    string okConfig = "raw:route[0]\n";
    string errConfig = "raw:"
        "route[1]\n"
        "route[0].name \"foo\"\n"
        "route[0].selector \"foo bar\"\n"
        "route[0].feed \"baz\"\n";
    {
        DocumentProtocol protocol(_loadTypes, _repo, okConfig);
        EXPECT_TRUE(dynamic_cast<DocumentRouteSelectorPolicy*>(protocol.createPolicy("DocumentRouteSelector", "").get()) != NULL);
        EXPECT_TRUE(dynamic_cast<ErrorPolicy*>(protocol.createPolicy("DocumentRouteSelector", errConfig).get()) != NULL);
    }
    {
        DocumentProtocol protocol(_loadTypes, _repo, errConfig);
        EXPECT_TRUE(dynamic_cast<ErrorPolicy*>(protocol.createPolicy("DocumentRouteSelector", "").get()) != NULL);
        EXPECT_TRUE(dynamic_cast<DocumentRouteSelectorPolicy*>(protocol.createPolicy("DocumentRouteSelector", okConfig).get()) != NULL);
    }

    // Test policy with proper config.
    TestFrame frame(_repo);
    frame.setHop(mbus::HopSpec("test", "[DocumentRouteSelector:raw:"
                               "route[2]\n"
                               "route[0].name \"foo\"\n"
                               "route[0].selector \"testdoc\"\n"
                               "route[0].feed \"myfeed\"\n"
                               "route[1].name \"bar\"\n"
                               "route[1].selector \"other\"\n"
                               "route[1].feed \"myfeed\"\n]")
                 .addRecipient("foo")
                 .addRecipient("bar"));

    frame.setMessage(mbus::Message::UP(new GetDocumentMessage(document::DocumentId("doc:scheme:"), 0)));
    EXPECT_TRUE(frame.testSelect(StringList().add("foo").add("bar")));

    mbus::Message::UP put(new PutDocumentMessage(Document::SP(
                            new Document(*_docType,
                                    DocumentId("doc:scheme:")))));
    frame.setMessage(std::move(put));
    EXPECT_TRUE(frame.testSelect( StringList().add("foo")));

    {
        vdslib::OperationList opList;

        document::DocumentId id("doc:scheme:");
        Document::UP doc(new Document(*_docType, id));
        opList.addPut(std::move(doc));

        document::BucketIdFactory factory;
        put = frame.setMessage(MultiOperationMessage::create(_repo, factory.getBucketId(id), opList));
        EXPECT_TRUE(frame.testSelect(StringList().add("foo")));
    }

    {
        vdslib::OperationList opList;
        document::DocumentId id("doc:scheme:");
        Document::UP doc(new Document(*_repo->getDocumentType("other"), id));
        opList.addPut(std::move(doc));

        document::BucketIdFactory factory;
        put = frame.setMessage(MultiOperationMessage::create(_repo,
            factory.getBucketId(id), opList));
        EXPECT_TRUE(frame.testSelect(StringList().add("bar")));
    }

    frame.setMessage(mbus::Message::UP(new RemoveDocumentMessage(document::DocumentId("doc:scheme:"))));
    EXPECT_TRUE(frame.testSelect(StringList().add("foo").add("bar")));

    frame.setMessage(mbus::Message::UP(new UpdateDocumentMessage(
                            document::DocumentUpdate::SP(
                                    new document::DocumentUpdate(
                                            *_docType,
                                            DocumentId("doc:scheme:"))))));
    EXPECT_TRUE(frame.testSelect(StringList().add("foo")));

    frame.setMessage(std::move(put));
    EXPECT_TRUE(frame.testMergeOneReply("foo"));
}

void
Test::testDocumentRouteSelectorIgnore()
{
    TestFrame frame(_repo);
    frame.setHop(mbus::HopSpec("test", "[DocumentRouteSelector:raw:"
                               "route[1]\n"
                               "route[0].name \"docproc/cluster.foo\"\n"
                               "route[0].selector \"testdoc and testdoc.stringfield == 'foo'\"\n"
                               "route[0].feed \"myfeed\"\n]")
                 .addRecipient("docproc/cluster.foo"));

    frame.setMessage(mbus::Message::UP(new PutDocumentMessage(
                            document::Document::SP(
                                    new document::Document(*_docType,
                                            DocumentId("id:yarn:testdoc:n=1234:fluff"))))));
    std::vector<mbus::RoutingNode*> leaf;
    ASSERT_TRUE(frame.select(leaf, 0));
    mbus::Reply::UP reply = frame.getReceptor().getReply(600);
    ASSERT_TRUE(reply.get() != NULL);
    EXPECT_EQUAL(uint32_t(DocumentProtocol::REPLY_DOCUMENTIGNORED), reply->getType());
    EXPECT_EQUAL(0u, reply->getNumErrors());

    frame.setMessage(mbus::Message::UP(new UpdateDocumentMessage(
                            document::DocumentUpdate::SP(
                                    new document::DocumentUpdate(
                                            *_docType,
                                            DocumentId("doc:scheme:"))))));
    EXPECT_TRUE(frame.testSelect(StringList().add("docproc/cluster.foo")));
}

namespace {
    string getDefaultDistributionConfig(
                    uint16_t redundancy = 2, uint16_t nodeCount = 10,
                    storage::lib::Distribution::DiskDistribution distr
                            = storage::lib::Distribution::MODULO_BID)
    {
        std::ostringstream ost;
        ost << "raw:redundancy " << redundancy << "\n"
            << "group[1]\n"
            << "group[0].index \"invalid\"\n"
            << "group[0].name \"invalid\"\n"
            << "group[0].partitions \"*\"\n"
            << "group[0].nodes[" << nodeCount << "]\n";
        for (uint16_t i=0; i<nodeCount; ++i) {
            ost << "group[0].nodes[" << i << "].index " << i << "\n";
        }
        ost << "disk_distribution "
            << storage::lib::Distribution::getDiskDistributionName(distr)
            << "\n";
        return ost.str();
    }
}

void Test::testLoadBalancer() {
    LoadBalancer lb("foo", "");

    IMirrorAPI::SpecList entries;
    entries.push_back(IMirrorAPI::Spec("foo/0/default", "tcp/bar:1"));
    entries.push_back(IMirrorAPI::Spec("foo/1/default", "tcp/bar:2"));
    entries.push_back(IMirrorAPI::Spec("foo/2/default", "tcp/bar:3"));

    const std::vector<LoadBalancer::NodeInfo>& nodeInfo = lb.getNodeInfo();

    for (int i = 0; i < 99; i++) {
        std::pair<string, int> recipient = lb.getRecipient(entries);
        EXPECT_EQUAL((i % 3), recipient.second);
    }

    // Simulate that one node is overloaded. It returns busy twice as often as the others.
    for (int i = 0; i < 100; i++) {
        lb.received(0, true);
        lb.received(0, false);
        lb.received(0, false);
        lb.received(2, true);
        lb.received(2, false);
        lb.received(2, false);
        lb.received(1, true);
        lb.received(1, true);
        lb.received(1, false);
    }

    EXPECT_EQUAL(421, (int)(100 * nodeInfo[0].weight / nodeInfo[1].weight));
    EXPECT_EQUAL(421, (int)(100 * nodeInfo[2].weight / nodeInfo[1].weight));

    EXPECT_EQUAL(0 , lb.getRecipient(entries).second);
    EXPECT_EQUAL(0 , lb.getRecipient(entries).second);
    EXPECT_EQUAL(1 , lb.getRecipient(entries).second);
    EXPECT_EQUAL(2 , lb.getRecipient(entries).second);
    EXPECT_EQUAL(2 , lb.getRecipient(entries).second);
    EXPECT_EQUAL(2 , lb.getRecipient(entries).second);
    EXPECT_EQUAL(2 , lb.getRecipient(entries).second);
    EXPECT_EQUAL(0 , lb.getRecipient(entries).second);
    EXPECT_EQUAL(0 , lb.getRecipient(entries).second);
    EXPECT_EQUAL(0 , lb.getRecipient(entries).second);
}

void
Test::requireThatStoragePolicyWithIllegalParamIsAnErrorPolicy()
{
    EXPECT_TRUE(isErrorPolicy("Storage", ""));
    EXPECT_TRUE(isErrorPolicy("Storage", "config=foo;slobroks=foo"));
    EXPECT_TRUE(isErrorPolicy("Storage", "slobroks=foo"));
}

void
Test::requireThatStoragePolicyIsRandomWithoutState()
{
    TestFrame frame(_repo);
    frame.setMessage(newPutDocumentMessage("doc:scheme:"));

    mbus::Slobrok slobrok;
    std::vector<mbus::TestServer*> servers;
    for (uint32_t i = 0; i < 5; ++i) {
        mbus::TestServer *srv = new mbus::TestServer(
                mbus::Identity(vespalib::make_string("storage/cluster.mycluster/distributor/%d", i)),
                mbus::RoutingSpec(), slobrok, "",
                mbus::IProtocol::SP(new DocumentProtocol(_loadTypes, _repo)));
        servers.push_back(srv);
        srv->net.registerSession("default");
    }
    string param = vespalib::make_string(
            "cluster=mycluster;slobroks=tcp/localhost:%d;clusterconfigid=%s;syncinit",
            slobrok.port(), getDefaultDistributionConfig(2, 5).c_str());
    StoragePolicy &policy = setupStoragePolicy(
            frame, param,
            "storage/cluster.mycluster/distributor/*/default", 5);
    ASSERT_TRUE(policy.getSystemState() == NULL);

    std::set<string> lst;
    for (uint32_t i = 0; i < 666; i++) {
        std::vector<mbus::RoutingNode*> leaf;
        ASSERT_TRUE(frame.select(leaf, 1));
        lst.insert(leaf[0]->getRoute().toString());
        leaf[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));
    }
    EXPECT_EQUAL(servers.size(), lst.size());
    for (uint32_t i = 0; i < servers.size(); ++i) {
        delete servers[i];
    }
}

StoragePolicy &
Test::setupStoragePolicy(TestFrame &frame, const string &param,
                         const string &pattern, int32_t numEntries)
{
    frame.setHop(mbus::HopSpec("test", vespalib::make_string("[Storage:%s]", param.c_str())));
    mbus::MessageBus &mbus = frame.getMessageBus();
    const mbus::HopBlueprint *hop = mbus.getRoutingTable(DocumentProtocol::NAME)->getHop("test");
    const mbus::PolicyDirective dir = static_cast<mbus::PolicyDirective&>(*hop->getDirective(0));
    StoragePolicy &policy = static_cast<StoragePolicy&>(*mbus.getRoutingPolicy(
                    DocumentProtocol::NAME,
                    dir.getName(),
                    dir.getParam()));
    policy.initSynchronous();
    assertMirrorReady(*policy.getMirror());
    if (numEntries >= 0) {
        assertMirrorContains(*policy.getMirror(), pattern, numEntries);
    }
    return policy;
}

void
Test::requireThatStoragePolicyIsTargetedWithState()
{
    TestFrame frame(_repo);
    frame.setMessage(newPutDocumentMessage("doc:scheme:"));

    mbus::Slobrok slobrok;
    std::vector<mbus::TestServer*> servers;
    for (uint32_t i = 0; i < 5; ++i) {
        mbus::TestServer *srv = new mbus::TestServer(
                mbus::Identity(vespalib::make_string("storage/cluster.mycluster/distributor/%d", i)),
                mbus::RoutingSpec(), slobrok, "",
                mbus::IProtocol::SP(new DocumentProtocol(_loadTypes, _repo)));
        servers.push_back(srv);
        srv->net.registerSession("default");
    }
    string param = vespalib::make_string(
            "cluster=mycluster;slobroks=tcp/localhost:%d;clusterconfigid=%s;syncinit",
            slobrok.port(), getDefaultDistributionConfig(2, 5).c_str());
    StoragePolicy &policy = setupStoragePolicy(
            frame, param,
            "storage/cluster.mycluster/distributor/*/default", 5);
    ASSERT_TRUE(policy.getSystemState() == NULL);
    {
        std::vector<mbus::RoutingNode*> leaf;
        ASSERT_TRUE(frame.select(leaf, 1));
        leaf[0]->handleReply(mbus::Reply::UP(new WrongDistributionReply("distributor:5 storage:5")));
        ASSERT_TRUE(policy.getSystemState() != NULL);
        EXPECT_EQUAL(policy.getSystemState()->toString(), "distributor:5 storage:5");
    }
    std::set<string> lst;
    for (int i = 0; i < 666; i++) {
        std::vector<mbus::RoutingNode*> leaf;
        ASSERT_TRUE(frame.select(leaf, 1));
        lst.insert(leaf[0]->getRoute().toString());
        leaf[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));
    }
    EXPECT_EQUAL(1u, lst.size());
    for (uint32_t i = 0; i < servers.size(); ++i) {
        delete servers[i];
    }
}

void
Test::requireThatStoragePolicyCombinesSystemAndSlobrokState()
{
    TestFrame frame(_repo);
    frame.setMessage(newPutDocumentMessage("doc:scheme:"));

    mbus::Slobrok slobrok;
    mbus::TestServer server(mbus::Identity("storage/cluster.mycluster/distributor/0"),
                            mbus::RoutingSpec(), slobrok, "",
                            mbus::IProtocol::SP(new DocumentProtocol(_loadTypes, _repo)));
    server.net.registerSession("default");

    string param = vespalib::make_string(
            "cluster=mycluster;slobroks=tcp/localhost:%d;clusterconfigid=%s;syncinit",
            slobrok.port(), getDefaultDistributionConfig(2, 5).c_str());
    StoragePolicy &policy = setupStoragePolicy(
            frame, param,
            "storage/cluster.mycluster/distributor/*/default", 1);
    ASSERT_TRUE(policy.getSystemState() == NULL);
    {
        std::vector<mbus::RoutingNode*> leaf;
        ASSERT_TRUE(frame.select(leaf, 1));
        leaf[0]->handleReply(mbus::Reply::UP(new WrongDistributionReply("distributor:99 storage:99")));
        ASSERT_TRUE(policy.getSystemState() != NULL);
        EXPECT_EQUAL(policy.getSystemState()->toString(), "distributor:99 storage:99");
    }
    for (int i = 0; i < 666; i++) {
        ASSERT_TRUE(frame.testSelect(StringList().add(server.net.getConnectionSpec() + "/default")));
    }
}

void
Test::testSubsetService()
{
    // Prepare message.
    TestFrame frame(_repo, "docproc/cluster.default");
    frame.setMessage(mbus::Message::UP(new PutDocumentMessage(Document::SP(
                                    new Document(*_docType,
                                            DocumentId("doc:scheme:"))))));

    // Test requerying for adding nodes.
    frame.setHop(mbus::HopSpec("test", "docproc/cluster.default/[SubsetService:2]/chain.default"));
    std::set<string> lst;
    for (uint32_t i = 1; i <= 10; ++i) {
        frame.getNetwork().registerSession(vespalib::make_string("%d/chain.default", i));
        ASSERT_TRUE(frame.waitSlobrok("docproc/cluster.default/*/chain.default", i));

        std::vector<mbus::RoutingNode*> leaf;
        ASSERT_TRUE(frame.select(leaf, 1));
        lst.insert(leaf[0]->getRoute().toString());
        leaf[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));
        ASSERT_TRUE(frame.getReceptor().getReply(600).get() != NULL);
    }
    ASSERT_TRUE(lst.size() > 1); // must have requeried

    // Test load balancing.
    string prev = "";
    for (uint32_t i = 1; i <= 10; ++i) {
        std::vector<mbus::RoutingNode*> leaf;
        ASSERT_TRUE(frame.select(leaf, 1));

        string next = leaf[0]->getRoute().toString();
        if (prev.empty()) {
            ASSERT_TRUE(!next.empty());
        } else {
            ASSERT_TRUE(prev != next);
        }

        prev = next;
        leaf[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));
        ASSERT_TRUE(frame.getReceptor().getReply(600).get() != NULL);
    }

    // Test requerying for dropping nodes.
    lst.clear();
    for (uint32_t i = 1; i <= 10; ++i) {
        std::vector<mbus::RoutingNode*> leaf;
        ASSERT_TRUE(frame.select(leaf, 1));
        string route = leaf[0]->getRoute().toString();
        lst.insert(route);

        frame.getNetwork().unregisterSession(route.substr(frame.getIdentity().length() + 1));
        ASSERT_TRUE(frame.waitSlobrok("docproc/cluster.default/*/chain.default", 10 - i));

        mbus::Reply::UP reply(new mbus::EmptyReply());
        reply->addError(mbus::Error(mbus::ErrorCode::NO_ADDRESS_FOR_SERVICE, route));
        leaf[0]->handleReply(std::move(reply));
        ASSERT_TRUE(frame.getReceptor().getReply(600).get() != NULL);
    }
    EXPECT_EQUAL(10u, lst.size());

    // Test merge behavior.
    frame.setHop(mbus::HopSpec("test", "[SubsetService]"));
    EXPECT_TRUE(frame.testMergeOneReply("*"));
}

void
Test::testSubsetServiceCache()
{
    TestFrame fooFrame(_repo, "docproc/cluster.default");
    mbus::HopSpec fooHop("foo", "docproc/cluster.default/[SubsetService:2]/chain.foo");
    fooFrame.setMessage(mbus::Message::UP(new GetDocumentMessage(document::DocumentId("doc:scheme:foo"))));
    fooFrame.setHop(fooHop);

    TestFrame barFrame(fooFrame);
    mbus::HopSpec barHop("bar", "docproc/cluster.default/[SubsetService:2]/chain.bar");
    barFrame.setMessage(mbus::Message::UP(new GetDocumentMessage(document::DocumentId("doc:scheme:bar"))));
    barFrame.setHop(barHop);

    fooFrame.getMessageBus().setupRouting(
            mbus::RoutingSpec().addTable(mbus::RoutingTableSpec(DocumentProtocol::NAME)
                                         .addHop(fooHop)
                                         .addHop(barHop)));

    fooFrame.getNetwork().registerSession("0/chain.foo");
    fooFrame.getNetwork().registerSession("0/chain.bar");
    ASSERT_TRUE(fooFrame.waitSlobrok("docproc/cluster.default/0/*", 2));

    std::vector<mbus::RoutingNode*> fooSelected;
    ASSERT_TRUE(fooFrame.select(fooSelected, 1));
    EXPECT_EQUAL("docproc/cluster.default/0/chain.foo", fooSelected[0]->getRoute().getHop(0).toString());

    std::vector<mbus::RoutingNode*> barSelected;
    ASSERT_TRUE(barFrame.select(barSelected, 1));
    EXPECT_EQUAL("docproc/cluster.default/0/chain.bar", barSelected[0]->getRoute().getHop(0).toString());

    barSelected[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));
    fooSelected[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));

    ASSERT_TRUE(barFrame.getReceptor().getReply(600).get() != NULL);
    ASSERT_TRUE(fooFrame.getReceptor().getReply(600).get() != NULL);
}

bool
Test::trySelect(TestFrame &frame, uint32_t numSelects, const std::vector<string> &expected) {
    std::set<string> lst;
    for (uint32_t i = 0; i < numSelects; ++i) {
        std::vector<mbus::RoutingNode*> leaf;
        if (!expected.empty()) {
            frame.select(leaf, 1);
            lst.insert(leaf[0]->getRoute().toString());
            leaf[0]->handleReply(mbus::Reply::UP(new mbus::EmptyReply()));
        } else {
            frame.select(leaf, 0);
        }
        if(frame.getReceptor().getReply(600).get() == NULL) {
            LOG(error, "Reply failed to propagate to reply handler.");
            return false;
        }
    }
    if (expected.size() != lst.size()) {
        LOG(error, "Expected %d recipients, got %d.", (uint32_t)expected.size(), (uint32_t)lst.size());
        return false;
    }
    std::set<string>::iterator it = lst.begin();
    for (uint32_t i = 0; i < expected.size(); ++i, ++it) {
        if (*it != expected[i]) {
            LOG(error, "Expected '%s', got '%s'.", expected[i].c_str(), it->c_str());
            return false;
        }
    }
    return true;
}

bool
Test::isErrorPolicy(const string &name, const string &param)
{
    DocumentProtocol protocol(_loadTypes, _repo);
    mbus::IRoutingPolicy::UP policy = protocol.createPolicy(name, param);

    return policy.get() != NULL && dynamic_cast<ErrorPolicy*>(policy.get()) != NULL;
}

